{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_LANES = 2\n",
    "INPUT_SIZE = (288, 800)\n",
    "CLS_SHAPE = (20, 100, NUM_LANES)\n",
    "\n",
    "INPUT_SHAPE = (*INPUT_SIZE, 3)\n",
    "OUTPUT_SHAPE = CLS_SHAPE\n",
    "\n",
    "PREFETCH_SIZE = 200\n",
    "\n",
    "DTYPE = tf.float32\n",
    "\n",
    "TAKE_LLAMAS_MAX = 1000\n",
    "TAKE_CULANE_MAX = 0\n",
    "TAKE_MIWULA_MAX = 0 # all of them\n",
    "\n",
    "# Ratio of validation data vs training data\n",
    "VALID_RATIO = 0.3\n",
    "\n",
    "BATCH_SIZE = 8 # tune to available GPU memory\n",
    "PREFETCH_SIZE = 200 # tune to available memory\n",
    "\n",
    "\n",
    "EPOCHS = 20 # training epochs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display(display_list):\n",
    "    plt.figure(figsize=(15, 15))\n",
    "    \n",
    "    title = ['Input Image', 'True Mask', 'Predicted Mask']\n",
    "\n",
    "    for i in range(len(display_list)):\n",
    "        plt.subplot(1, len(display_list), i+1)\n",
    "        plt.title(title[i])\n",
    "        plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))\n",
    "        plt.axis('off')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from unsupervised_llamas.label_scripts.spline_creator import get_horizontal_values_for_four_lanes\n",
    "\n",
    "BASE_PATH = '/Volumes/tank_data/Archiv/ML/llamas/labels/'\n",
    "\n",
    "LLAMAS_SHAPE = (717, 1276, 3)\n",
    "#MASK_SHAPE = (712, 1272, 3)\n",
    "\n",
    "IMAGE_SHAPE = INPUT_SHAPE\n",
    "MASK_SHAPE = OUTPUT_SHAPE\n",
    "\n",
    "PARALLEL_JOBS = 4# tf.data.experimental.AUTOTUNE\n",
    "\n",
    "\n",
    "def resize_img(img, shape):\n",
    "    w = img.shape[1]\n",
    "    h = img.shape[0]\n",
    "    \n",
    "    ratio = shape[1] / shape[0]\n",
    "    \n",
    "    tgt = img[h-int(w/ratio):h,0:w]\n",
    "    tgt = cv2.resize(tgt, (shape[1], shape[0]))\n",
    "    return tgt\n",
    "\n",
    "def draw_spline(img, points, color):\n",
    "    n = np.array(points)\n",
    "    idx = np.arange(0, len(n)).reshape(-1,1)\n",
    "    idx = idx[n >= 0]\n",
    "    n = n[n >= 0]\n",
    "    n = n.reshape(-1,1)\n",
    "    pts = np.hstack((n, idx))\n",
    "    \n",
    "    cv2.polylines(img, np.int32([pts]), isClosed=False, color=color, thickness=10, lineType=cv2.LINE_8)\n",
    "    \n",
    "    del n\n",
    "    del idx\n",
    "    del pts\n",
    "\n",
    "def generate_mask(file):\n",
    "    file = file.numpy().decode(\"utf-8\")\n",
    "    mask = np.zeros(LLAMAS_SHAPE, dtype='uint8')\n",
    "    \n",
    "    lines = get_horizontal_values_for_four_lanes(file)\n",
    "    l0 = lines[1]\n",
    "    r0 = lines[2]\n",
    "    \n",
    "    draw_spline(mask, l0, (0,255,0))\n",
    "    draw_spline(mask, r0, (255,0,0))\n",
    "    \n",
    "    mask = resize_img(mask, MASK_SHAPE)\n",
    "    return mask[:,:,0:2]\n",
    "\n",
    "def read_image(file):\n",
    "    file = file.numpy().decode(\"utf-8\")\n",
    "    with open(file, 'r') as fp:\n",
    "        meta = json.load(fp)\n",
    "        \n",
    "    image_path = os.path.dirname(file)\n",
    "    image_path = image_path.replace('/labels/', '/color_images/')\n",
    "    img_name = meta['image_name'] + '_color_rect.png'\n",
    "    fname = os.path.join(image_path, img_name)\n",
    "    \n",
    "    img = cv2.imread(fname)\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "    return resize_img(img, IMAGE_SHAPE)\n",
    "\n",
    "def process_json(json):\n",
    "    img = tf.py_function(read_image, [json], Tout=DTYPE)\n",
    "    mask = tf.py_function(generate_mask, [json], Tout=DTYPE)\n",
    "    \n",
    "    img = tf.reshape(img, shape=IMAGE_SHAPE)\n",
    "    mask = tf.reshape(mask, shape=MASK_SHAPE)\n",
    "    return img, mask\n",
    "\n",
    "SHUFFLE_SIZE = 100000\n",
    "\n",
    "llamas_train_ds = tf.data.Dataset.list_files(os.path.join(BASE_PATH, 'train', '*/*.json'))\n",
    "llamas_train_ds = llamas_train_ds.shuffle(SHUFFLE_SIZE)\n",
    "\n",
    "llamas_train_ds = llamas_train_ds.take(TAKE_LLAMAS_MAX)\n",
    "LLAMAS_TRAIN_LEN = len(list(llamas_train_ds))\n",
    "print('LLAMAS training images: %d' % LLAMAS_TRAIN_LEN)\n",
    "\n",
    "llamas_train_ds = llamas_train_ds.map(process_json, num_parallel_calls=PARALLEL_JOBS)\n",
    "\n",
    "llamas_valid_ds = tf.data.Dataset.list_files(os.path.join(BASE_PATH, 'valid', '*/*.json'))\n",
    "llamas_valid_ds = llamas_valid_ds.shuffle(SHUFFLE_SIZE)\n",
    "\n",
    "llamas_valid_ds = llamas_valid_ds.take(TAKE_LLAMAS_MAX)\n",
    "LLAMAS_VALID_LEN = len(list(llamas_valid_ds))\n",
    "print('LLAMAS validation images: %d' % LLAMAS_VALID_LEN)\n",
    "\n",
    "llamas_valid_ds = llamas_valid_ds.map(process_json, num_parallel_calls=PARALLEL_JOBS)\n",
    "\n",
    "for img, mask in llamas_valid_ds.take(3):\n",
    "    m = tf.zeros((20,100,1), dtype=DTYPE)\n",
    "    m = tf.concat([mask, m], -1)\n",
    "    display([img, m])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tflite_runtime.interpreter as tflite\n",
    "\n",
    "interpreter = tflite.Interpreter('ultrafast_quant_edgetpu.tflite', experimental_delegates=[tflite.load_delegate('libedgetpu.1.dylib')])\n",
    "interpreter.allocate_tensors()\n",
    "\n",
    "def set_input(interpreter, tensor):\n",
    "    input_details = interpreter.get_input_details()[0]\n",
    "    #print(input_details)\n",
    "    tensor_index = input_details['index']\n",
    "    input_tensor = interpreter.tensor(tensor_index)()[0]\n",
    "    \n",
    "    scale, zero_point = input_details['quantization']\n",
    "    \n",
    "    input_tensor[:, :] = tensor\n",
    "    \n",
    "def run_ultrafast(interpreter, inp_img):\n",
    "    set_input(interpreter, inp_img)\n",
    "    \n",
    "    #s = time.time()\n",
    "    interpreter.invoke()\n",
    "    #e = time.time()\n",
    "    #print(e-s)\n",
    "    \n",
    "    output_details = interpreter.get_output_details()[0]\n",
    "    #print(output_details)\n",
    "    output = interpreter.get_tensor(output_details['index'])\n",
    "    # Outputs from the TFLite model are uint8, so we dequantize the results:\n",
    "    #scale, zero_point = output_details['quantization']\n",
    "    #output = scale * (output - zero_point)\n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llamas_valid_ds = llamas_valid_ds.prefetch(100).cache()\n",
    "\n",
    "for img, mask in llamas_valid_ds.take(50):\n",
    "    output = run_ultrafast(interpreter, img)\n",
    "    \n",
    "    m = tf.zeros((20,100,1), dtype=DTYPE)\n",
    "    m = tf.concat([output[0], m], -1)\n",
    "    \n",
    "    m2 = tf.zeros((20,100,1), dtype=DTYPE)\n",
    "    m2 = tf.concat([mask, m2], -1)\n",
    "    \n",
    "    display([img, m2, m])\n",
    "    \n",
    "    #print(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def resize_img(img, shape):\n",
    "    w = img.shape[1]\n",
    "    h = img.shape[0]\n",
    "    \n",
    "    ratio = shape[1] / shape[0]\n",
    "    \n",
    "    tgt = img[h-int(w/ratio):h,0:w]\n",
    "    tgt = cv2.resize(tgt, (shape[1], shape[0]))\n",
    "    return tgt\n",
    "\n",
    "\n",
    "def pipeline(img):\n",
    "    img = resize_img(img, INPUT_SHAPE)\n",
    "    \n",
    "    out = run_ultrafast(interpreter, img)\n",
    "    m = tf.zeros((20,100,1), dtype=DTYPE)\n",
    "    m = tf.concat([out[0], m], -1)\n",
    "    \n",
    "    #print(m.shape)\n",
    "    #out_img = (m.numpy() * 255).astype(np.int)\n",
    "    #print(out_img)\n",
    "    ratio = m.shape[0] / m.shape[1]\n",
    "    \n",
    "    \n",
    "    new_w = img.shape[1]\n",
    "    new_h = int(new_w * ratio)\n",
    "    \n",
    "    rescale = cv2.resize(m.numpy(), (new_w, new_h))\n",
    "    \n",
    "    min_val = np.min(rescale)\n",
    "    max_val = np.max(rescale)\n",
    "    rescale = ((rescale - min_val) / (max_val - min_val)) * 255\n",
    "    \n",
    "    tmp = np.zeros(img.shape, dtype=np.uint8)\n",
    "    tmp[128:288,:] = rescale\n",
    "    \n",
    "    overlay = cv2.addWeighted(img, 1, tmp, 1, 0)\n",
    "    \n",
    "    return overlay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = cv2.imread('test.png')\n",
    "test = cv2.cvtColor(test, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "\n",
    "out = pipeline(test)\n",
    "\n",
    "display([test, out])\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from moviepy.editor import VideoFileClip\n",
    "\n",
    "clip = 'out2.mp4'\n",
    "\n",
    "clip1 = VideoFileClip(clip)\n",
    "driving_clip = clip1.fl_image(pipeline)\n",
    "driving_clip.write_videofile('out2_ultrafast.mp4', audio=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
